{-# OPTIONS -fno-warn-missing-signatures #-}
{-# LANGUAGE CPP, ScopedTypeVariables, TypeOperators, FlexibleContexts, TypeFamilies #-}
{- |
Module      :  Fitting
Description :  <optional short text displayed on contents page>
Copyright   :  (c) Sergey Sichevskiy 2012
License     :  BSD3

Maintainer  :  s.sichevskij@gmal.com
Stability   :  experimental
Portability :  non-portable (use the non-portable module 'Numeric.LevMar')

Определение параметров модели по исходным данным

-}

module Fitting where

import Helper
import Types
import Model
import Annealing            ( annealing )

import Data.List
import Data.HList
import Data.Maybe
import Data.Array.IO
import Data.Packed.Matrix   ( Matrix )
import Data.Packed.Vector   ( toList, fromList )

import Control.Monad
import Control.Applicative  ( (<$>) )
import Control.Monad.Reader ( asks, liftIO )
import Control.Exception    ( evaluate )
import Control.Concurrent
--import Control.Parallel

import System.Time
import System.CPUTime
import System.IO.Unsafe

import Text.Printf         (printf)

import Numeric.LevMar      ( levmar, defaultOpts, Constraints(..), LevMarError, Info )


--fittedTo :: Model [AnyPhotometry] -> [AnyPhotometry] -> Factory [ModelParameter]
--fittedTo :: ( Model xs ys ) => (xs -> ys) -> (ys,xs) -> Factory (xs, Info Double)
fittedTo :: ( HFromList Double ps
            , HMap ToDouble xs
            , HMap ToDouble ys
            , HToList (Map ToDouble xs) Double
            , HToList (Map ToDouble ys) Double
            )
          =>  (ps -> xs, xs -> ys) -> (ys,xs) -> Factory (Either LevMarError (xs, Info Double, Matrix Double))
fittedTo (m,f) (d,p) = do

  i <- asks $ maxIterations . fitmethod
--  a <- asks $ absTolerance  . fitmethod
--  r <- asks $ relTolerance  . fitmethod

  let r = levmar f' Nothing p' d' i o c

      f' = fromList <$> (parEval v) . hToList . hMap ToDouble . f . m . hFromList . toList
--      f' = fromList <$> parEval . hToList . hMap ToDouble . f . m . hFromList . toList
--      f' = fromList <$> hToList . hMap ToDouble . f . m . hFromList . toList
      p' = fromList  $  hToList $ hMap ToDouble p
      d' = fromList  $  hToList $ hMap ToDouble d

      o  = defaultOpts -- { optStopNorm2Dp = 0, optStopNormInfJacTe = 0 }

      v  = unsafePerformIO $ listMVar 4

      c  = Constraints { lowerBounds       = Nothing
                       , upperBounds       = Nothing
                       , weights           = Nothing
                       , linearConstraints = Nothing
                       }


  case r of
    Right (ps,info,mt) -> return $ Right (m $ hFromList $ toList ps, info, mt)
    Left  err          -> return $ Left err

fittedByASTo (m,f) (d,p) = do

--  i <- asks $ maxIterations . fitmethod
--  a <- asks $ absTolerance  . fitmethod
--  r <- asks $ relTolerance  . fitmethod

  let f' = cost d' <$> (parEval v) . hToList . hMap ToDouble . f . m
--      f' = fromList <$> hToList . hMap ToDouble . f . m . hFromList . toList
      p' = hMap ToDouble p
      d' = hToList $ hMap ToDouble d

      v  = unsafePerformIO $ listMVar 4

      cost :: [Double] -> [Double] -> Double
      cost ds ys = sum $ map (**2) $ zipWith (-) ys ds

  (x,e) <- liftIO $ annealing f' p' 1 50000 1e-8

  return $ (m x, e)

fitTo :: ( xs' ~ Zip xs (Map ToDouble xs)
         , xs  ~ Map FromDouble xs'
         , HZip xs (Map ToDouble xs)         
         , HMap FromDouble xs'
         , HMap ToDouble xs
         , HMap ToDouble ys
         , HFromList Double (Map ToDouble xs)
         , HToList (Map ToDouble xs) Double
         , HToList (Map ToDouble ys) Double
         )
       =>  (xs -> ys) -> (ys,xs,(xs,xs)) -> Factory (Either LevMarError (xs, Info Double, Matrix Double))
fitTo m (d,p,b) = do

  i <- asks $ maxIterations . fitmethod
--  a <- asks $ absTolerance  . fitmethod
--  r <- asks $ relTolerance  . fitmethod

  let r = levmar f' Nothing p' d' i o c

     -- f' = fromList <$> (parEval v) . (m $$ p) . toList
      f' = fromList <$> (m $$ p) . toList
      p' = fromList  $  hToList $ hMap ToDouble p
      d' = fromList  $  hToList $ hMap ToDouble d

      o  = defaultOpts -- { optStopNorm2Dp = 0, optStopNormInfJacTe = 0 }

     -- v  = unsafePerformIO $ listMVar 4

      c  = Constraints { lowerBounds       = Just $ fromList $ hToList $ hMap ToDouble ( fst b )
                       , upperBounds       = Just $ fromList $ hToList $ hMap ToDouble ( snd b )
                       , weights           = Nothing
                       , linearConstraints = Nothing
                       }

  case r of
    Right (ps,info,mt) -> return $ Right (p `updateBy` (toList ps), info, mt)
    Left  err          -> return $ Left err 

#ifdef PAR_EVAL

parEval :: [(MVar Double, MVar Double)] -> [Double] -> [Double]
parEval inpOut = unsafeDupablePerformIO.parEvalIO inpOut

{--
parEval :: [Double] -> [Double]
parEval     [] = []
parEval (x:xs) = x `par` es `pseq` (x : es)
  where
    es = parEval xs

parEval a (x:xs) = unsafePerformIO $ do

                mv <- newEmptyMVar

                forkIO $ do
                  e <- evaluate x

                  --putStrLn $ "thread : "++ show e

                  putMVar mv e

                es <- evaluate $ parEval a xs
                e  <- takeMVar mv

                return (x:xs)
--}

parEvalIO :: [(MVar Double, MVar Double)] -> [Double] -> IO [Double]
parEvalIO _      [] = return []
parEvalIO inpOut xs = do

  let (inp,out) = unzip inpOut
      n = min (length inp) (length xs)

  zipWithM_ putMVar inp xs

  es <- mapM takeMVar $ take n out

  ees <- parEvalIO inpOut $ drop n xs

  return $ es ++ ees

listMVar :: Int -> IO [(MVar Double, MVar Double)]
listMVar n = do
  inp <- replicateM n newEmptyMVar
  out <- replicateM n newEmptyMVar

  let inpOut = zip inp out

  mapM_ (forkIO . evalFunc) inpOut

  return inpOut

  where
    evalFunc (inp,out) = takeMVar inp >>= evaluate >>= putMVar out >> evalFunc (inp,out)

#else

parEval :: (Chan (Int,Double), Chan (Int,Double)) -> [Double] -> [Double]
parEval inpOut = unsafeDupablePerformIO.parEvalIO inpOut

parEvalIO :: (Chan (Int,a), Chan (Int,a)) -> [a] -> IO [a]
parEvalIO _      [] = return []
parEvalIO (inp,out) xs = do

  writeList2Chan inp $ zip [1..] xs

--  start <- getCPUTime

--  t0 <- getClockTime
  
  es <- (map snd) <$> (sortBy sortFunc) <$> replicateM n (readChan out)

--  end   <- getCPUTime

--  let diff = (fromIntegral (end - start)) / (10^12)

  es' <- mapM evaluate es

--  putStrLn ("total sum: " ++ show es)

--  t1 <- getClockTime

--  printf "total time : t0 : %17.6f seconds\ntotal time : t1 : %17.6f seconds\ntotal time :      %17.6f seconds\n"  (sec t0) (sec t1) (secDiff t0 t1)

--  putStrLn ("total time: " ++ show diff ++ " seconds")

  return es'

  where
     n = length xs

--     sortFunc :: (Int,Double) -> (Int,Double) -> Ordering
     sortFunc (a,_) (b,_) = a `compare` b

listMVar :: Int -> IO (Chan (Int,a), Chan (Int,a))
listMVar n = do

  inp <- newChan
  out <- newChan

  mapM_ (forkIO . evalFunc inp ) $ replicate n out

  return (inp,out)

  where
    evalFunc inp out = do

        (i,x) <- readChan inp

--        start <- getCPUTime 

--        t0 <- getClockTime

        e <- evaluate x

--        end   <- getCPUTime

--        t1 <- getClockTime

        writeChan out  (i,e)

--        let diff = (fromIntegral (end - start)) / (10^12)  

--        putStrLn ("s: " ++ show (i,e))

--        printf "t : t0 : %17.6f seconds\nt : t1 : %17.6f seconds\nt :      %17.6f seconds\n" (sec t0) (sec t1) (secDiff t0 t1)

--        putStrLn ("t: " ++ show diff ++ " seconds")

        evalFunc inp out
--}

secDiff :: ClockTime -> ClockTime -> Double
secDiff (TOD secs1 psecs1) (TOD secs2 psecs2) = fromInteger (psecs2 - psecs1) / 1e12 + fromInteger (secs2 - secs1)

sec :: ClockTime -> Double
sec (TOD secs psecs) = fromInteger secs + (fromInteger psecs/1e12)

#endif